import tkinter as tk
from tkinter import filedialog, messagebox
import xarray as xr
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import cartopy.crs as ccrs
import numpy as np
import cartopy
import sys
import os

# 处理打包后的资源路径
if getattr(sys, 'frozen', False):
    base_path = sys._MEIPASS
    cartopy.config['data_dir'] = os.path.join(base_path, 'cartopy', 'data')

# 确保中文显示
plt.rcParams["font.family"] =  ["SimSun"]

# 全局变量
ds = None
current_canvas = None
plottable_vars = []  # 存储可绘制的变量


def open_file():
    global ds, current_canvas, plottable_vars

    # 清除现有图表
    if current_canvas:
        current_canvas.get_tk_widget().destroy()

    file_path = filedialog.askopenfilename()
    if not file_path:
        return  # 用户取消选择

    try:
        # 打开文件并更新全局变量
        ds = xr.open_dataset(file_path, engine="netcdf4")

        # 分析变量类型并筛选可绘制的变量
        plottable_vars = []
        var_names = list(ds.variables.keys())

        for name in var_names:
            var = ds[name]
            dims = var.dims
            shape = var.shape

            # 判断变量类型
            var_type = None
            if len(dims) == 1:
                var_type = "1D"
            else:
                # 检查可能的纬度和经度变量名
                lat_vars = ["lat", "Latitude"]
                lon_vars = ["lon", "Longitude"]
                is_geo2d = False
                for lat_var in lat_vars:
                    for lon_var in lon_vars:
                        if lat_var in ds and lon_var in ds:
                            lat_dims = ds[lat_var].dims[0]
                            lon_dims = ds[lon_var].dims[0]
                            if lat_dims in dims and lon_dims in dims:
                                is_geo2d = True
                                break
                    if is_geo2d:
                        break
                if is_geo2d:
                    var_type = "Geo2D"
                elif len(dims) == 2:
                    var_type = "2D"
                else:
                    continue

            # 仅添加可绘制的变量
            if var_type in ["1D", "2D", "Geo2D"]:
                plottable_vars.append((name, var_type))

        # 显示变量列表（带类型）
        var_list.delete(0, tk.END)
        for name, var_type in plottable_vars:
            var_list.insert(tk.END, f"{name} ({var_type})" )

        # 启用绘图按钮
        btn_plot.config(state=tk.NORMAL)

        status_label.config(text=f"已加载文件: {file_path}")

    except Exception as e:
        error_msg = f"无法打开文件: {file_path}\n错误详情: {str(e)}"
        print(error_msg)
        messagebox.showerror("错误", error_msg)
        ds = None
        btn_plot.config(state=tk.DISABLED)


def plot_variable():
    global current_canvas

    if ds is None:
        error_msg = "错误: 请先打开一个文件"
        print(error_msg)
        messagebox.showerror("错误", error_msg)
        return

    selected = var_list.curselection()
    if not selected:
        error_msg = "提示: 请选择一个变量"
        print(error_msg)
        messagebox.showinfo("提示", error_msg)
        return

    # 获取选中的变量名（去除类型后缀）
    selected_text = var_list.get(selected[0])
    var_name = selected_text.split(" (")[0]
    var_type = selected_text.split(" (")[1].rstrip(")")

    try:
        var = ds[var_name]

        # 清除 matplotlib 中的旧图形
        plt.clf()

        # 创建新图表，将图形与指定的画布关联
        fig = plt.figure(figsize=(10, 6))

        def process_high_dimensional_var(current_var):
            # 获取非地理维度（非经纬度维度）
            lat_vars = ["lat", "Latitude"]
            lon_vars = ["lon", "Longitude"]
            lat_var = next((v for v in lat_vars if v in ds), None)
            lon_var = next((v for v in lon_vars if v in ds), None)
            geo_dims = set()
            if lat_var:
                geo_dims.add(ds[lat_var].dims[0])
            if lon_var:
                geo_dims.add(ds[lon_var].dims[0])
            extra_dims = [dim for dim in current_var.dims if dim not in geo_dims]

            # 如果没有额外维度，返回当前变量
            if not extra_dims:
                return current_var

            # 创建弹窗选择维度和索引
            select_window = tk.Toplevel(root)
            select_window.title("选择维度和索引")

            # 选择维度
            dim_var = tk.StringVar(select_window)
            dim_var.set(extra_dims[0])
            dim_menu = tk.OptionMenu(select_window, dim_var, *extra_dims)
            dim_menu.pack(padx=10, pady=5)

            # 选择索引
            dim_index = current_var.dims.index(dim_var.get())
            index_options = list(range(current_var.shape[dim_index]))
            index_var = tk.StringVar(select_window)
            index_var.set(str(index_options[0]))
            index_menu = tk.OptionMenu(select_window, index_var, *index_options)
            index_menu.pack(padx=10, pady=5)

            # 当维度改变时，更新索引选项
            def update_index_options(*args):
                dim_index = current_var.dims.index(dim_var.get())
                index_options = list(range(current_var.shape[dim_index]))
                menu = index_menu['menu']
                menu.delete(0, 'end')
                for option in index_options:
                    menu.add_command(label=option, command=tk._setit(index_var, option))
                index_var.set(str(index_options[0]))

            dim_var.trace_add('write', update_index_options)

            def select_dim_and_index(): 
                nonlocal current_var
                selected_dim = dim_var.get()
                selected_index = int(index_var.get())
                current_var = current_var.isel({selected_dim: selected_index})
                select_window.destroy()
                # 递归处理剩余高维变量
                if len(current_var.dims) > 2:
                    current_var = process_high_dimensional_var(current_var)
                proceed_plotting(current_var)

            btn_select = tk.Button(select_window, text="选择", command=select_dim_and_index)
            btn_select.pack(padx=10, pady=10)
            select_window.wait_window()

        def proceed_plotting(processed_var):
            # 创建新的弹窗用于显示图表
            plot_window = tk.Toplevel(root)
            plot_window.title(f"{var_name} 绘图")
            plot_window.geometry("1000x800")

            # 根据变量类型绘图
            if var_type == "1D":
                processed_var.plot()
                plt.title(f"{var_name} - 折线图")

            elif var_type in ["2D", "Geo2D"]:
                if var_type == "Geo2D":
                    ax = fig.add_subplot(111, projection=ccrs.PlateCarree())
                    
                    lat_vars = ["lat", "Latitude"]
                    lon_vars = ["lon", "Longitude"]
                    lat_var = next((v for v in lat_vars if v in ds), None)
                    lon_var = next((v for v in lon_vars if v in ds), None)

                    if lat_var and lon_var:
                        im = ax.pcolormesh(
                            ds[lon_var],
                            ds[lat_var],
                            processed_var,
                            transform=ccrs.PlateCarree(),
                            cmap="viridis"
                        )
                        ax.coastlines()
                        ax.gridlines(draw_labels=True)
                        plt.colorbar(im, label=processed_var.attrs.get("units", "Unknown"))
                        plt.title(f"{var_name} - 地理分布图")
                    else:
                        error_msg = "错误: 未找到可用的经纬度变量"
                        print(error_msg)
                        messagebox.showerror("错误", error_msg)
                        return
                else:
                    processed_var.plot()
                    plt.title(f"{var_name} - 热力图")
            else:
                error_msg = f"提示: 变量 {var_name} 类型不支持绘图 ({var_type})" 
                print(error_msg)
                messagebox.showinfo("提示", error_msg)
                return

            # 在新弹窗中显示图表
            canvas = FigureCanvasTkAgg(fig, master=plot_window)
            canvas.draw()
            canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
            
        # 处理高维变量
        if len(var.dims) > 2:
            process_high_dimensional_var(var)
            return

        # 原有的二维及以下变量绘图逻辑
        proceed_plotting(var)

    except Exception as e:
        error_msg = f"绘图错误: 无法绘制变量 {var_name}: {e}"
        print(error_msg)
        messagebox.showerror("绘图错误", error_msg)


# 创建Tkinter界面
root = tk.Tk()
root.title("HDF/NetCDF数据可视化工具")
root.geometry("800x600")

# 创建顶部按钮区域
button_frame = tk.Frame(root)
button_frame.pack(fill=tk.X, padx=10, pady=10)

btn_open = tk.Button(button_frame, text="打开文件", command=open_file)
btn_open.pack(side=tk.LEFT, padx=5)

btn_plot = tk.Button(button_frame, text="绘制选中变量", command=plot_variable, state=tk.DISABLED)
btn_plot.pack(side=tk.LEFT, padx=5)

# 状态标签
status_label = tk.Label(root, text="请打开一个HDF/NetCDF文件", fg="gray")
status_label.pack(pady=5)

# 创建变量列表区域
list_frame = tk.Frame(root)
list_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

tk.Label(list_frame, text="可用变量（仅显示可绘制类型）:").pack(anchor=tk.W)

var_list = tk.Listbox(list_frame, selectmode=tk.SINGLE)
var_list.pack(fill=tk.BOTH, expand=True)

# 添加滚动条
scrollbar = tk.Scrollbar(var_list, orient=tk.VERTICAL)
scrollbar.config(command=var_list.yview)
scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
var_list.config(yscrollcommand=scrollbar.set)

# 图表区域（动态创建）
chart_frame = tk.Frame(root)
chart_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

root.mainloop()